import os
import re
import argparse
import json
from pathlib import Path


def extract_code_blocks(text):
    """
    从DeepSeek返回文本中提取所有代码块
    返回: 字典 {语言: 代码内容}
    """
    code_blocks = {}
    pattern = r"```(\w+)?\n(.*?)```"
    matches = re.finditer(pattern, text, re.DOTALL)

    for i, match in enumerate(matches):
        lang = match.group(1) or f"block_{i + 1}"
        code = match.group(2).strip()
        code_blocks[lang] = code

    return code_blocks


def update_source_files(code_blocks, output_dir="."):
    """
    将代码块写入对应文件
    文件名映射规则:
      1. 语言标记为文件名 (如 ```python → file.py)
      2. 特殊标记映射 (如 ```src/main.py → main.py)
      3. 默认保存为block_X.txt
    """
    file_mapping = {
        "python": ".py",
        "javascript": ".js",
        "html": ".html",
        "css": ".css",
        "java": ".java",
        "cpp": ".cpp"
    }

    for lang, code in code_blocks.items():
        # 检查是否是路径格式 (包含/)
        if '/' in lang or '\\' in lang:
            file_path = Path(lang)
        elif lang in file_mapping:
            file_path = Path(f"{lang}{file_mapping[lang]}")
        else:
            file_path = Path(f"block_{lang}.txt")

        # 创建目录
        full_path = Path(output_dir) / file_path
        full_path.parent.mkdir(parents=True, exist_ok=True)

        # 写入文件
        with open(full_path, "w", encoding="utf-8") as f:
            f.write(code)
        print(f"✅ 代码已写入: {full_path}")


def process_deepseek_output(input_file, output_dir):
    """处理DeepSeek输出文件并更新源码"""
    with open(input_file, "r", encoding="utf-8") as f:
        content = f.read()

    code_blocks = extract_code_blocks(content)

    if not code_blocks:
        print("⚠️ 未检测到代码块")
        return

    update_source_files(code_blocks, output_dir)

def process_deepseek_text(content, output_dir):
    """处理DeepSeek输出文件并更新源码"""
    code_blocks = extract_code_blocks(content)

    if not code_blocks:
        print("⚠️ 未检测到代码块")
        return

    update_source_files(code_blocks, output_dir)

def main():
    parser = argparse.ArgumentParser(description="DeepSeek代码自动更新工具")
    parser.add_argument("-i", "--input", required=True, help="DeepSeek生成的文本文件")
    parser.add_argument("-o", "--output", default="output", help="输出目录(默认: output)")
    args = parser.parse_args()

    process_deepseek_output(args.input, args.output)


if __name__ == "__main__":
    main()